#include <iostream>
#include <vector>
#include <time.h>

// DBoW2
#include "DBoW2.h"

// OpenCV
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/features2d.hpp>

using namespace DBoW2;
using namespace std;

const int NIMAGES = 4;

//command line parser
class CmdLineParser {
  int argc;
  char **argv;
public:
  CmdLineParser(int _argc, char **_argv) : argc(_argc), argv(_argv) {}

  bool operator[](string param) {
    int idx = -1;
    for (int i = 0; i < argc && idx == -1; i++) if (string(argv[i]) == param) idx = i;
    return (idx != -1);
  }

  string operator()(string param, string defvalue = "-1") {
    int idx = -1;
    for (int i = 0; i < argc && idx == -1; i++) if (string(argv[i]) == param) idx = i;
    if (idx == -1) return defvalue; else return (argv[idx + 1]);
  }
};


void changeStructure(const cv::Mat &plain, vector<cv::Mat> &out) {
  out.resize(plain.rows);

  for (int i = 0; i < plain.rows; ++i) {
    out[i] = plain.row(i);
  }
}

void loadFeatures(vector<vector<cv::Mat>> &features) {
  features.clear();
  features.reserve(NIMAGES);

  cv::Ptr<cv::ORB> orb = cv::ORB::create();

  cout << "Extracting ORB features..." << endl;
  for (int i = 0; i < NIMAGES; ++i) {
    stringstream ss;
    ss << "images/image" << i << ".png";

    cv::Mat image = cv::imread(ss.str(), 0);
    cv::Mat mask;
    vector<cv::KeyPoint> keypoints;
    cv::Mat descriptors;

    orb->detectAndCompute(image, mask, keypoints, descriptors);

    features.push_back(vector<cv::Mat>());
    changeStructure(descriptors, features.back());
  }
}

// ----------------------------------------------------------------------------

int main(int argc, char **argv) {

  try {
    CmdLineParser cml(argc, argv);
    if (cml["-h"] || argc != 3) {
      cerr << "Usage: ./test_io_binary in.yml out.dbow" << endl;
      return -1;
    }

    // Step1: extract orb features
    vector<vector<cv::Mat>> features;
    loadFeatures(features);

    // Step2: load voc from yaml file and test bow vector encoding
    cout << "\n==================yaml==================" << endl;
    // 2.1 load voc from yaml file
    OrbVocabulary voc;
    const clock_t begin_time = clock();
    voc.load(argv[1]);
    float seconds = float(clock( ) - begin_time) / 1000;
    std::cout << "INFO=yaml||loadTime=" << seconds << "ms" << std::endl;

    cout << "INFO=yaml format voc file loaded." << endl;

    // 2.2 test bow vector encoding
    BowVector v1, v2;
    for(int i = 0; i < NIMAGES; i++)
    {
      voc.transform(features[i], v1);
      for(int j = 0; j < NIMAGES; j++)
      {
        voc.transform(features[j], v2);

        double score = voc.score(v1, v2);
        cout << "INFO=yaml||tag=Image " << i << " vs Image " << j << ": " << score << endl;
      }
    }

    // Step3: load voc from binary file and test bow vector encoding
    cout << "\n==================binary==================" << endl;
    // 3.1 save voc to binary file
    voc.save(argv[2]);
    cout << "INFO=binary format voc file saved." << endl;

    // 3.2 load voc from binary file
    const clock_t begin_time1 = clock();
    voc.load(argv[2]);
    float seconds1 = float(clock( ) - begin_time1) / 1000;
    std::cout << "INFO=binary||loadTime=" << seconds1 << "ms" << std::endl;
    cout << "INFO=binary format voc file loaded." << endl;

    // 3.3 test bow vector encoding
    for(int i = 0; i < NIMAGES; i++)
    {
      voc.transform(features[i], v1);
      for(int j = 0; j < NIMAGES; j++)
      {
        voc.transform(features[j], v2);

        double score = voc.score(v1, v2);
        cout << "INFO=binary||tag=Image " << i << " vs Image " << j << ": " << score << endl;
      }
    }

  } catch (std::exception &ex) {
    cerr << ex.what() << endl;
  }

  return 0;
}
